#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from typing import Any, Generic, TypeVar, Type
from uuid import UUID
from pydantic import BaseModel
from sqlmodel import SQLModel, select, func, Session
from sqlmodel.ext.asyncio.session import AsyncSession
from sqlmodel.sql.expression import Select
from sqlalchemy import exc
from enum import Enum

from common.exception.errors import CustomException
from common.response.response_code import StandardResponseCode
from .db import SessionLocal, get_db

ModelType = TypeVar("ModelType", bound=SQLModel)
CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)
SchemaType = TypeVar("SchemaType", bound=BaseModel)
T = TypeVar("T", bound=SQLModel)


class IOrderEnum(str, Enum):
    ascendent = "ascendent"
    descendent = "descendent"


class CRUDBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
    def __init__(self, model: type[ModelType]):
        """
        CRUD object with default methods to Create, Read, Update, Delete (CRUD).
        
        Args
        * `model`: A SQLModel model class
        * `schema`: A Pydantic model (schema) class
        """
        self.model = model

    async def get(
        self, *, id: int, db_session: AsyncSession | None = None
    ) -> ModelType | None:
        db_session = db_session or SessionLocal()
        async with db_session as session:
            query = select(self.model).where(self.model.id == id)
            response = await session.exec(query)
            return response.one_or_none()

    async def get_by_ids(
        self,
        *,
        list_ids: list[UUID | str],
        db_session: AsyncSession | None = None,
    ) -> list[ModelType] | None:
        db_session = db_session or SessionLocal()
        async with db_session as session:
            response = await session.exec(
                select(self.model).where(self.model.id.in_(list_ids))
            )
            return response.all()

    async def count(
        self, db_session: AsyncSession | None = None
    ) -> ModelType | None:
        db_session = db_session or SessionLocal()
        async with db_session as session:
            response = await session.exec(
                select(func.count()).select_from(select(self.model).subquery())
            )
            return response.one_or_none()

    async def page(
        self,
        *,
        page: int = 0,
        page_size: int = 100,
        order_by: str | None = None,
        order: IOrderEnum | None = IOrderEnum.ascendent,
        db_session: AsyncSession | None = None,
    ) -> list[ModelType]:
        db_session = db_session or SessionLocal()

        columns = self.model.__table__.columns

        if order_by is None or order_by not in columns:
            order_by = "id"

        if order == IOrderEnum.ascendent:
            query = (
                select(self.model)
                .offset(page)
                .limit(page_size)
                .order_by(columns[order_by].asc())
            )
        else:
            query = (
                select(self.model)
                .offset(page)
                .limit(page_size)
                .order_by(columns[order_by].desc())
            )
        async with db_session as session:
            response = await session.exec(query)
            return response.all()

    async def create(
        self,
        *,
        obj_in: CreateSchemaType | ModelType,
        created_by_id: UUID | str | None = None,
        db_session: AsyncSession | None = None,
    ) -> ModelType:
        db_session = db_session or SessionLocal()
        db_obj = self.model.model_validate(obj_in)  # type: ignore

        if created_by_id:
            db_obj.created_by_id = created_by_id

        async with db_session as session:
            try:
                session.add(db_obj)
                await session.commit()
            except exc.IntegrityError:
                session.rollback()
                raise CustomException(
                    code=StandardResponseCode.HTTP_409,
                    msg="Resource already exists",
                )
            await session.refresh(db_obj)
            return db_obj

    async def update(
        self,
        *,
        obj_current: ModelType,
        obj_new: UpdateSchemaType | dict[str, Any] | ModelType,
        db_session: AsyncSession | None = None,
    ) -> ModelType:
        db_session = db_session or SessionLocal()

        if isinstance(obj_new, dict):
            update_data = obj_new
        else:
            update_data = obj_new.model_dump(
                exclude_unset=True
            )  # This tells Pydantic to not include the values that were not sent
        for field in update_data:
            setattr(obj_current, field, update_data[field])

        async with db_session as session:
            session.add(obj_current)
            await session.commit()
            await session.refresh(obj_current)
            return obj_current

    async def remove(
        self, *, id: UUID | str, db_session: AsyncSession | None = None
    ) -> ModelType:
        db_session = db_session or SessionLocal()
        async with db_session as session:
            response = await session.exec(
                select(self.model).where(self.model.id == id)
            )
            obj = response.one()
            await session.delete(obj)
            await session.commit()
            return obj